import os
import cv2
import json
import random
from tqdm import tqdm



def make_json():
    name_list = [p for p in os.listdir('../suichang_round1_train_210120\suichang_round1_train_210120') if 'tif' in p]
    random.shuffle(name_list)
    train_name_list = name_list[:int(0.7*len(name_list))]
    val_name_list = name_list[int(0.7*len(name_list)):]

    with open('label_file/train.json','w') as f:
        json.dump(train_name_list,f)
    with open('label_file/val.json','w') as f:
        json.dump(val_name_list,f)


def get_data_info(data_path):
    info = [0 for i in range(10)]
    for name in tqdm(os.listdir(data_path)):
        if 'png' in name:
            label_dir = os.path.join(data_path,name)
            label = cv2.imread(label_dir) - 1
            for i in range(10):
                info[i] += ((label==i).sum() / 10000)
    new_info = [p / max(info) for p in info]
    print(new_info)


if __name__ == '__main__':
    # make_json()
    get_data_info(data_path='../suichang_round1_train_210120\suichang_round1_train_210120')